import torch
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
from typing import Dict, Any
import torchvision.models as models


class ModelInterpreter:
    """模型解释器"""
    
    def __init__(self, model: torch.nn.Module, save_dir: Path):
        self.model = model
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(parents=True, exist_ok=True)
        
    def count_parameters(self) -> Dict[str, int]:
        """统计模型参数"""
        total_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        
        return {
            'total_parameters': total_params,
            'trainable_parameters': trainable_params,
            'non_trainable_parameters': total_params - trainable_params
        }
    
    def analyze_model_complexity(self) -> Dict[str, Any]:
        """分析模型复杂度"""
        param_stats = self.count_parameters()
        
        # 计算模型大小 (MB)
        param_size = sum(p.numel() * p.element_size() for p in self.model.parameters())
        buffer_size = sum(b.numel() * b.element_size() for b in self.model.buffers())
        model_size_mb = (param_size + buffer_size) / (1024 * 1024)
        
        # 统计层类型
        layer_types = {}
        for name, module in self.model.named_modules():
            module_type = type(module).__name__
            if module_type not in layer_types:
                layer_types[module_type] = 0
            layer_types[module_type] += 1
        
        return {
            **param_stats,
            'model_size_mb': model_size_mb,
            'layer_types': layer_types
        }
    
    def visualize_model_architecture(self) -> Dict[str, Any]:
        """可视化模型架构"""
        complexity_info = self.analyze_model_complexity()
        
        # 创建参数分布图
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        fig.suptitle('Model Architecture Analysis', fontsize=16)
        
        # 参数统计饼图
        param_labels = ['Trainable', 'Non-trainable']
        param_sizes = [complexity_info['trainable_parameters'], 
                      complexity_info['non_trainable_parameters']]
        
        axes[0, 0].pie(param_sizes, labels=param_labels, autopct='%1.1f%%')
        axes[0, 0].set_title('Parameter Distribution')
        
        # 层类型分布
        layer_types = complexity_info['layer_types']
        layer_names = list(layer_types.keys())[:10]  # 显示前10种类型
        layer_counts = [layer_types[name] for name in layer_names]
        
        axes[0, 1].bar(range(len(layer_names)), layer_counts)
        axes[0, 1].set_title('Layer Type Distribution')
        axes[0, 1].set_xticks(range(len(layer_names)))
        axes[0, 1].set_xticklabels(layer_names, rotation=45)
        
        # 参数量按层分布
        layer_params = []
        layer_labels = []
        
        for name, module in self.model.named_modules():
            if len(list(module.children())) == 0:  # 叶子节点
                param_count = sum(p.numel() for p in module.parameters())
                if param_count > 0:
                    layer_params.append(param_count)
                    layer_labels.append(name.split('.')[-1][:10])  # 截断长名称
        
        # 选择参数最多的前15层
        if len(layer_params) > 15:
            sorted_indices = np.argsort(layer_params)[-15:]
            layer_params = [layer_params[i] for i in sorted_indices]
            layer_labels = [layer_labels[i] for i in sorted_indices]
        
        axes[1, 0].barh(range(len(layer_params)), layer_params)
        axes[1, 0].set_title('Parameters by Layer (Top 15)')
        axes[1, 0].set_yticks(range(len(layer_labels)))
        axes[1, 0].set_yticklabels(layer_labels)
        
        # 模型信息文本
        info_text = f"""
Model Complexity Summary:
• Total Parameters: {complexity_info['total_parameters']:,}
• Trainable Parameters: {complexity_info['trainable_parameters']:,}
• Model Size: {complexity_info['model_size_mb']:.2f} MB
• Number of Layer Types: {len(complexity_info['layer_types'])}
        """
        
        axes[1, 1].text(0.1, 0.5, info_text, transform=axes[1, 1].transAxes,
                        fontsize=10, verticalalignment='center',
                        bbox=dict(boxstyle='round', facecolor='lightgray'))
        axes[1, 1].set_xlim(0, 1)
        axes[1, 1].set_ylim(0, 1)
        axes[1, 1].axis('off')
        axes[1, 1].set_title('Model Summary')
        
        plt.tight_layout()
        save_path = self.save_dir / 'model_architecture.png'
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"模型架构分析图已保存: {save_path}")
        return complexity_info
